import numpy as np
import faiss
from util import createXY
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
import argparse
import logging
from tqdm import tqdm
from FaissKNeighbors import FaissKNeighbors

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def get_args():
    parser = argparse.ArgumentParser(description='使用CPU或GPU训练KNN模型。')
    parser.add_argument('-m', '--mode', type=str, required=True, choices=['cpu', 'gpu'], help='训练模式：CPU或GPU')
    parser.add_argument('-f', '--feature', type=str, required=True, choices=['flat', 'vgg'], help='特征提取方法：flat或vgg')
    parser.add_argument('-l', '--library', type=str, required=True, choices=['sklearn', 'faiss'], help='使用的库：sklearn或faiss')
    return parser.parse_args()

def main():
    args = get_args()
    res = faiss.StandardGpuResources() if args.mode == 'gpu' else None
    logging.info(f"训练模式：{args.mode.upper()}")
    logging.info(f"特征提取：{args.feature.upper()}")
    logging.info(f"使用库：{args.library.upper()}")

    train_folder = r"D:\cat_dog_data\data\train"
    dest_folder = "."
    try:
        X, y = createXY(train_folder=train_folder, dest_folder=dest_folder, method=args.feature)
    except Exception as e:
        logging.error(f"特征和标签生成失败：{str(e)}")
        return

    X = np.array(X).astype('float32')
    if len(X.shape) != 2:
        logging.error(f"X数据形状异常，当前形状：{X.shape}")
        return
    faiss.normalize_L2(X)
    y = np.array(y).astype('int32')
    logging.info(f"数据加载完成 | 特征维度：{X.shape} | 标签数量：{y.shape}")

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.25, random_state=2023, stratify=y
    )
    logging.info(f"数据集划分完成 | 训练集：{X_train.shape} | 测试集：{X_test.shape}")

    best_k, best_accuracy = -1, 0.0
    k_values = range(1, 6)
    KNNClass = FaissKNeighbors if args.library == 'faiss' else KNeighborsClassifier

    for k in tqdm(k_values, desc='寻找最佳k值'):
        if args.library == 'faiss':
            knn = KNNClass(k=k, res=res)
        else:
            knn = KNNClass(n_neighbors=k, n_jobs=-1)
        knn.fit(X_train, y_train)
        accuracy = knn.score(X_test, y_test)
        if accuracy > best_accuracy:
            best_accuracy = accuracy
            best_k = k

    logging.info(f'最佳k值: {best_k}, 最高准确率: {best_accuracy:.4f}')

if __name__ == '__main__':
    main()